from flask_script import Manager, Command, Option
from flask_migrate import Migrate, MigrateCommand

from main import app
from models import db
from models.user import User
from models.host import Host, HostGroup


class SuperAdmin(Command):
    "创建超级管理员"

    option_list = (
        Option('--username', '-u', dest='username'),
        Option('--password', '-p', dest='password'),
    )

    def run(self, username, password):
        try:
            admin_user = User(username=username)
            admin_user.set_password(password)
            db.session.add(admin_user)
            db.session.commit()
        except Exception as e:
            db.session.rollback()
            print("ERROR: 创建失败")
            print(e)


manager = Manager(app)
Migrate(app=app, db=db)
manager.add_command('db', MigrateCommand)   # 创建数据库映射命令
manager.add_command('create-admin', SuperAdmin())     # 创建管理员

if __name__ == '__main__':
    manager.run()
